"""
数字人服务器
"""

import asyncio
import json
import logging
from typing import Dict, Set, Optional, Any
import websockets
from websockets.server import WebSocketServerProtocol
from websockets.exceptions import ConnectionClosed, ConnectionClosedError

from .websocket_media import MediaStreamer, connection_manager
from .basereal import BaseReal

logger = logging.getLogger(__name__)

class DigitalHumanServer:
    """数字人服务器"""

    def __init__(self, host: str = "0.0.0.0", port: int = 8888):
        self.host = host
        self.port = port
        self.any4dh_reals: Dict[int, BaseReal] = {}
        self.server = None

        # 连接统计
        self.total_connections = 0
        self.active_connections = 0

    async def register_session(self, sessionid: int, any4dh_real: BaseReal) -> MediaStreamer:
        """注册新的数字人会话"""
        self.any4dh_reals[sessionid] = any4dh_real
        streamer = connection_manager.create_streamer(sessionid)
        return streamer

    async def unregister_session(self, sessionid: int):
        """注销数字人会话"""
        await connection_manager.remove_streamer(sessionid)
        if sessionid in self.any4dh_reals:
            del self.any4dh_reals[sessionid]

    async def handle_client(self, websocket: WebSocketServerProtocol):
        """处理客户端连接"""
        client_info = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        self.total_connections += 1
        self.active_connections += 1
        connection_manager.active_connections.add(websocket)

        try:
            async for message in websocket:
                await self.process_message(websocket, message)

        except ConnectionClosedError as e:
            logger.info(f"Client connection closed: {client_info}, code: {e.code}, reason: {e.reason}")
        except ConnectionClosed:
            logger.info(f"Client disconnected: {client_info}")
        except Exception as e:
            logger.error(f"Error handling client {client_info}: {e}")
        finally:
            self.active_connections -= 1
            connection_manager.active_connections.discard(websocket)

            # 清理会话连接
            await self.cleanup_client_sessions(websocket)

    async def cleanup_client_sessions(self, websocket: WebSocketServerProtocol):
        """清理客户端相关的会话"""
        for sessionid, streamer in list(connection_manager.streamers.items()):
            if streamer.websocket == websocket:
                await connection_manager.remove_streamer(sessionid)
                break

    async def process_message(self, websocket: WebSocketServerProtocol, message: str):
        """处理客户端消息"""
        try:
            data = json.loads(message)
            msg_type = data.get('type')

            if msg_type == 'join_session':
                await self.handle_join_session(websocket, data)
            elif msg_type == 'text_input':
                await self.handle_text_input(websocket, data)
            elif msg_type == 'audio_input':
                await self.handle_audio_input(websocket, data)
            elif msg_type == 'interrupt':
                await self.handle_interrupt(websocket, data)
            elif msg_type == 'ping':
                await self.handle_ping(websocket, data)
            elif msg_type == 'get_status':
                await self.handle_get_status(websocket, data)
            elif msg_type == 'set_audiotype':
                await self.handle_set_audiotype(websocket, data)
            elif msg_type == 'record_control':
                await self.handle_record_control(websocket, data)
            else:
                logger.warning(f"Unknown message type: {msg_type}")
                await self.send_error(websocket, f"Unknown message type: {msg_type}")

        except json.JSONDecodeError:
            logger.error("Invalid JSON message received")
            await self.send_error(websocket, "Invalid JSON message")
        except Exception as e:
            logger.error(f"Error processing message: {e}")
            await self.send_error(websocket, f"Message processing error: {str(e)}")

    async def handle_join_session(self, websocket: WebSocketServerProtocol, data: dict):
        """处理加入会话请求"""
        sessionid = data.get('sessionid')
        if sessionid and sessionid in connection_manager.streamers:
            streamer = connection_manager.streamers[sessionid]
            await streamer.register_client(websocket)

            # 发送加入成功消息
            await websocket.send(json.dumps({
                'type': 'session_joined',
                'sessionid': sessionid,
                'status': 'success',
                'timestamp': asyncio.get_event_loop().time()
            }))

            # 发送数字人状态
            if sessionid in self.any4dh_reals:
                any4dh_real = self.any4dh_reals[sessionid]
                await websocket.send(json.dumps({
                    'type': 'digital_human_status',
                    'sessionid': sessionid,
                    'is_speaking': any4dh_real.is_speaking(),
                    'timestamp': asyncio.get_event_loop().time()
                }))

        else:
            await self.send_error(websocket, 'Invalid session ID')

    async def handle_text_input(self, websocket: WebSocketServerProtocol, data: dict):
        """处理文本输入"""
        sessionid = data.get('sessionid')
        text = data.get('text', '')
        input_type = data.get('input_type', 'echo')  # 'echo' or 'chat'

        if sessionid and sessionid in self.any4dh_reals:
            any4dh_real = self.any4dh_reals[sessionid]

            try:
                if input_type == 'echo':
                    any4dh_real.put_msg_txt(text)
                elif input_type == 'chat':
                    # 异步调用LLM处理
                    from ...live_talking.llm import llm_response
                    asyncio.get_event_loop().run_in_executor(None, llm_response, text, any4dh_real)

                # 发送确认消息
                await websocket.send(json.dumps({
                    'type': 'text_input_received',
                    'sessionid': sessionid,
                    'text': text,
                    'input_type': input_type,
                    'timestamp': asyncio.get_event_loop().time()
                }))

            except Exception as e:
                logger.error(f"Error processing text input: {e}")
                await self.send_error(websocket, f"Text input processing failed: {str(e)}")
        else:
            await self.send_error(websocket, 'Session not found')

    async def handle_audio_input(self, websocket: WebSocketServerProtocol, data: dict):
        """处理音频输入"""
        sessionid = data.get('sessionid')
        audio_data = data.get('audio_data', '')

        if sessionid and sessionid in self.any4dh_reals:
            try:
                # 解码base64音频数据
                import base64
                audio_bytes = base64.b64decode(audio_data)
                self.any4dh_reals[sessionid].put_audio_file(audio_bytes)

                # 发送确认消息
                await websocket.send(json.dumps({
                    'type': 'audio_input_received',
                    'sessionid': sessionid,
                    'audio_size': len(audio_bytes),
                    'timestamp': asyncio.get_event_loop().time()
                }))

            except Exception as e:
                logger.error(f"Error processing audio input: {e}")
                await self.send_error(websocket, f"Audio input processing failed: {str(e)}")
        else:
            await self.send_error(websocket, 'Session not found')

    async def handle_interrupt(self, websocket: WebSocketServerProtocol, data: dict):
        """处理中断请求"""
        sessionid = data.get('sessionid')

        if sessionid and sessionid in self.any4dh_reals:
            try:
                self.any4dh_reals[sessionid].flush_talk()

                await websocket.send(json.dumps({
                    'type': 'interrupt_completed',
                    'sessionid': sessionid,
                    'timestamp': asyncio.get_event_loop().time()
                }))

            except Exception as e:
                logger.error(f"Error processing interrupt: {e}")
                await self.send_error(websocket, f"Interrupt processing failed: {str(e)}")
        else:
            await self.send_error(websocket, 'Session not found')

    async def handle_ping(self, websocket: WebSocketServerProtocol, data: dict):
        """处理ping消息"""
        await websocket.send(json.dumps({
            'type': 'pong',
            'timestamp': asyncio.get_event_loop().time(),
            'client_timestamp': data.get('timestamp')
        }))

    async def handle_get_status(self, websocket: WebSocketServerProtocol, data: dict):
        """处理状态查询"""
        sessionid = data.get('sessionid')

        if sessionid and sessionid in connection_manager.streamers:
            streamer = connection_manager.streamers[sessionid]
            stats = streamer.get_connection_stats()

            if sessionid in self.any4dh_reals:
                any4dh_real = self.any4dh_reals[sessionid]
                stats['is_speaking'] = any4dh_real.is_speaking()

            await websocket.send(json.dumps({
                'type': 'status_response',
                'sessionid': sessionid,
                'stats': stats,
                'timestamp': asyncio.get_event_loop().time()
            }))
        else:
            # 返回整体服务器状态
            await websocket.send(json.dumps({
                'type': 'server_status',
                'server_stats': {
                    'total_connections': self.total_connections,
                    'active_connections': self.active_connections,
                    'active_sessions': len(connection_manager.streamers)
                },
                'all_sessions_stats': connection_manager.get_all_stats(),
                'timestamp': asyncio.get_event_loop().time()
            }))

    async def handle_set_audiotype(self, websocket: WebSocketServerProtocol, data: dict):
        """处理音频类型设置"""
        sessionid = data.get('sessionid')
        audiotype = data.get('audiotype')
        reinit = data.get('reinit', True)

        if sessionid and sessionid in self.any4dh_reals:
            try:
                self.any4dh_reals[sessionid].set_custom_state(audiotype, reinit)

                await websocket.send(json.dumps({
                    'type': 'audiotype_set',
                    'sessionid': sessionid,
                    'audiotype': audiotype,
                    'reinit': reinit,
                    'timestamp': asyncio.get_event_loop().time()
                }))

            except Exception as e:
                logger.error(f"Error setting audiotype: {e}")
                await self.send_error(websocket, f"Set audiotype failed: {str(e)}")
        else:
            await self.send_error(websocket, 'Session not found')

    async def handle_record_control(self, websocket: WebSocketServerProtocol, data: dict):
        """处理录制控制"""
        sessionid = data.get('sessionid')
        record_type = data.get('record_type')  # 'start_record' or 'end_record'

        if sessionid and sessionid in self.any4dh_reals:
            try:
                if record_type == 'start_record':
                    self.any4dh_reals[sessionid].start_recording()
                elif record_type == 'end_record':
                    self.any4dh_reals[sessionid].stop_recording()
                else:
                    await self.send_error(websocket, 'Invalid record_type')
                    return

                await websocket.send(json.dumps({
                    'type': 'record_control_completed',
                    'sessionid': sessionid,
                    'record_type': record_type,
                    'timestamp': asyncio.get_event_loop().time()
                }))

            except Exception as e:
                logger.error(f"Error handling record control: {e}")
                await self.send_error(websocket, f"Record control failed: {str(e)}")
        else:
            await self.send_error(websocket, 'Session not found')

    async def send_error(self, websocket: WebSocketServerProtocol, error_message: str):
        """发送错误消息"""
        try:
            await websocket.send(json.dumps({
                'type': 'error',
                'error': error_message,
                'timestamp': asyncio.get_event_loop().time()
            }))
        except Exception as e:
            logger.error(f"Failed to send error message: {e}")

    async def start_server(self):
        """启动服务器"""

        self.server = await websockets.serve(
            self.handle_client,
            self.host,
            self.port,
            ping_interval=20,
            ping_timeout=10,
            close_timeout=10,
            max_size=10 * 1024 * 1024,  # 10MB max message size
            compression=None  # 禁用压缩，避免与图像压缩冲突
        )

        return self.server

    async def stop_server(self):
        """停止服务器"""
        if self.server:
            self.server.close()
            await self.server.wait_closed()

    async def cleanup_all_sessions(self):
        """清理所有会话"""
        for sessionid in list(connection_manager.streamers.keys()):
            await self.unregister_session(sessionid)

# 全局服务器实例
_server_instance: Optional[DigitalHumanServer] = None

def get_server() -> DigitalHumanServer:
    """获取全局服务器实例"""
    global _server_instance
    if _server_instance is None:
        _server_instance = DigitalHumanServer()
    return _server_instance

async def initialize_server(host: str = "0.0.0.0", port: int = 8888) -> DigitalHumanServer:
    """初始化服务器"""
    global _server_instance
    if _server_instance is None:
        _server_instance = DigitalHumanServer(host, port)
        await _server_instance.start_server()
    return _server_instance

async def shutdown_server():
    """关闭服务器"""
    global _server_instance
    if _server_instance:
        await _server_instance.stop_server()
        _server_instance = None